package mmq.mina.codec;

import java.io.IOException;

import org.apache.mina.core.buffer.IoBuffer;
import org.apache.mina.core.session.IoSession;
import org.apache.mina.filter.codec.CumulativeProtocolDecoder;
import org.apache.mina.filter.codec.ProtocolDecoderOutput;

import com.google.protobuf.CodedInputStream;
import com.google.protobuf.CodedOutputStream;
import com.google.protobuf.ExtensionRegistry;
import com.google.protobuf.Message;
import com.google.protobuf.GeneratedMessage.GeneratedExtension;

public class ProtobufDecoder extends CumulativeProtocolDecoder  {
	
	private final Message prototype;
	private final ExtensionRegistry extentions;
	
	public ProtobufDecoder(Message prototype, GeneratedExtension<? extends Message, ?>... extensions) {
		this.prototype = prototype;
		this.extentions = ExtensionRegistry.newInstance();
		for(GeneratedExtension<? extends Message, ?> extension : extensions)
			extentions.add(extension);
	}
	
	@Override
	protected boolean doDecode(IoSession session, IoBuffer in, ProtocolDecoderOutput out) throws Exception {
		SizeContext ctx = SizeContext.get(session, in);
		if(ctx.hasEnoughData(in)) {
			try	{
				Message.Builder builder = prototype.newBuilderForType();
				ctx.getInputStream(in).readMessage(builder, extentions);
				out.write(builder.build());
				return true;
			} finally {
				ctx.shiftPositionAndReset(session, in);
			}
		}
		return false;
	}
	
	static class SizeContext {
		
		private static String KEY = SizeContext.class.getCanonicalName();
		
		private final int size;
	
		static SizeContext get(IoSession session, IoBuffer in) throws IOException {
			SizeContext ctx = (SizeContext) session.getAttribute(KEY);
			if(ctx == null) {
				int size = CodedInputStream.newInstance(in.array(), in.position(), 5).readRawVarint32();
				ctx = new SizeContext(size, CodedOutputStream.computeRawVarint32Size(size));
				session.setAttribute(KEY, ctx);
			}
			return ctx;
		}
		
		private SizeContext(int messageLength, int sizeLenght) {
			this.size = messageLength + sizeLenght;
		}
		
		boolean hasEnoughData(IoBuffer in) {
			return in.remaining() >= size;
		}
		
		CodedInputStream getInputStream(IoBuffer in) {
			return CodedInputStream.newInstance(
				in.array(),
				in.position(),
				size);
		}
		
		void shiftPositionAndReset(IoSession session, IoBuffer in) {
			in.position(in.position() + size);
			session.removeAttribute(KEY);
		}
		
	}

}
